package com.box.ai.controller;

import com.box.ai.common.result.Result;
import com.box.ai.mapper.PDFUploadMapper;
import com.box.ai.pojo.entity.PdfChat;
import com.box.ai.pojo.vo.MessageVO;
import com.box.ai.service.PDFUploadService;
import com.box.ai.service.handler.ChatTypeHandler;
import com.box.ai.service.impl.FunctionTypeFactory;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@RequiredArgsConstructor
@RestController
@RequestMapping("/ai")
@Slf4j
public class TalkChatController {

    private final ChatClient gameChatClient;
    private final ChatClient chatClient;
    private final ChatClient customerClient;
    private final ChatClient pdfClient;
    private final FunctionTypeFactory functionTypeFactory;
    private final PDFUploadService pdfUploadService;
    private final PDFUploadMapper pdfUploadMapper;
    private final VectorStore redisVectorStore;

    /**
     * 智能对话功能实现
     * 
     * @param prompt
     * @param chatId
     * @return
     */
    @RequestMapping(value = "/chat", produces = "text/html;charset=utf-8")
    public Flux<String> chat(String prompt, String chatId) {
        // 会话之前保存会话id,通过工厂加策略模式
        ChatTypeHandler chatTypeHandler = functionTypeFactory.getBean("chat");
        chatTypeHandler.saveChat("chat", chatId);
        return chatClient.prompt()
                .user(prompt)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }

    /**
     * 游戏功能实现
     * 
     * @param prompt
     * @param chatId
     * @return
     */
    // produces限定用户发送的编码为utf8，Flux是Spring封装的一个响应式流式，大模型返回的内容会立即返回
    @RequestMapping(value = "/game", produces = "text/html;charset=utf-8")
    public Flux<String> game(String prompt, String chatId) {
        // 链式编程
        // prompt是构建提示词的过程，通|过对配置Bean中的system和用户发送的user信息构建提示词
        // advisors限定不同的会话id让会话不会再混乱（多个会话使用同一个记忆）
        // stream流式返回回答的过程
        // content构建大模型的返回内容
        ChatTypeHandler chatTypeHandler = functionTypeFactory.getBean("game");
        chatTypeHandler.saveChat("game", chatId);
        return gameChatClient.prompt()
                .user(prompt)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }

    /**
     * 智能客服功能实现
     * 
     * @param prompt
     * @param chatId
     * @return
     */
    // produces限定用户发送的编码为utf8，Flux是Spring封装的一个响应式流式，大模型返回的内容会立即返回
    @RequestMapping(value = "/service", produces = "text/html;charset=utf-8")
    public Flux<String> customer(String prompt, String chatId) {
        // 链式编程
        // prompt是构建提示词的过程，通|过对配置Bean中的system和用户发送的user信息构建提示词
        // advisors限定不同的会话id让会话不会再混乱（多个会话使用同一个记忆）
        // stream流式返回回答的过程
        // content构建大模型的返回内容
        ChatTypeHandler chatTypeHandler = functionTypeFactory.getBean("service");
        chatTypeHandler.saveChat("service", chatId);
        return customerClient.prompt()
                .user(prompt)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }

    /**
     * 基于PDF的智能问答
     * 
     * @param prompt
     * @param chatId
     * @return
     */
    @RequestMapping(value = "/pdf/chat", produces = "text/html;charset=utf-8")
    public Flux<String> chatPDF(String prompt, String chatId) {
        ChatTypeHandler chatTypeHandler = functionTypeFactory.getBean("pdf");
        chatTypeHandler.saveChat("pdf", chatId);
        String fileName = pdfUploadService.getFileName(chatId);


        return pdfClient.prompt()
                .user(prompt)
                .advisors(new QuestionAnswerAdvisor(redisVectorStore, SearchRequest
                        .builder()
                        .similarityThreshold(0.1d)
                        .filterExpression( "file_name == '"+fileName+"'")
                        .topK(6)
                        .build())
                )
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }

    /**
     * PDF上传
     * 
     * @param file
     * @param chatId
     * @return
     */
    @RequestMapping(value = "/pdf/upload/{chatId}")
    public Result upload(@RequestParam("file") MultipartFile file, @PathVariable String chatId) {
        try {
            // 1. 校验文件是否为PDF格式
            if (!Objects.equals(file.getContentType(), "application/pdf")) {
                return Result.fail("只能上传PDF文件！");
            }
            // 2.保存文件
            boolean success = pdfUploadService.upload(file, chatId);
            if (!success) {
                return Result.fail("保存文件失败！");
            }
            // 3.写入向量库
            this.writeToVectorStore(file.getResource());
            return Result.ok();
        } catch (Exception e) {
            log.error("Failed to upload PDF.", e);
            return Result.fail("上传文件失败！");
        }
    }

    /**
     * PDF下载
     * 
     * @param file
     * @param chatId
     * @return
     */
    /**
     * 文件下载
     */
    @GetMapping("/pdf/file/{chatId}")
    public ResponseEntity<Resource> download(@PathVariable("chatId") String chatId) throws IOException {
        // 1.读取文件
        Resource resource = pdfUploadService.getFile(chatId);
        if (!resource.exists()) {
            return ResponseEntity.notFound().build();
        }

        PdfChat pdfChat = pdfUploadMapper.selectByChatId(chatId);
        if (pdfChat == null || pdfChat.getFileName() == null) {
            return ResponseEntity.notFound().build();
        }

        // 2.文件名编码，写入响应头
        String filename = URLEncoder.encode(pdfChat.getFileName(), StandardCharsets.UTF_8);
        // 3.返回文件
        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .header("Content-Disposition", "attachment; filename=\"" + filename + "\"")
                .body(resource);
    }

    private void writeToVectorStore(Resource resource) {
        // 1.创建PDF的读取器
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource, // 文件源
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1) // 每1页PDF作为一个Document
                        .build());
        // 2.读取PDF文档，拆分为Document
        List<Document> documents = reader.read();
        // 3.为每个Document添加文件名元数据
        String fileName = resource.getFilename();
        documents.forEach(doc -> doc.getMetadata().put("file_name", fileName));
        // 4.写入向量库
        redisVectorStore.add(documents);
    }

    /**
     * 通过业务id查询所有会话
     * 
     * @param type
     * @return
     */
    @RequestMapping(value = "/history/{type}")
    public List<String> getChatIds(@PathVariable String type) {
        ChatTypeHandler chatTypeHandler = functionTypeFactory.getBean(type);
        List<String> typeList = chatTypeHandler.queryByType(type);
        return typeList;
    }

    /**
     * 查询会话id记录
     * 
     * @param type
     * @return
     */
    @RequestMapping(value = "/history/{type}/{chatId}")
    public List<MessageVO> save(@PathVariable String type, @PathVariable String chatId) {
        ChatTypeHandler chatTypeHandler = functionTypeFactory.getBean(type);
        List<MessageVO> messageVOS = chatTypeHandler.getHistoryChat(type, chatId);
        return messageVOS;
    }
}